#!/usr/bin/env python

##
## Copyright (C) 2009-2010 Alberto Ferrarini <alberto.ferrarini@univr.it>
## Contrib: Luciano Xumerle <luciano.xumerle@gmail.com>
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

import sys
import os.path
try:
        import pysqlite2.dbapi2 as sqlite
except:
        import sqlite3 as sqlite

errorMsg="SYNTAX:\n - " + sys.argv[0] + " <knownGene.txt|ensembl.gtf> <DataBase> [annotations files]\n"

if len(sys.argv)<3:
        print >>sys.stderr, errorMsg
        sys.exit()

### OPEN AND SET DATABASE
con = sqlite.connect(sys.argv[2])
cur = con.cursor()

def initdb():
	cur.execute("CREATE TABLE knowngene (id varchar,chr varchar,strand varchar,chr_start long,chr_stop long,type varchar)")
	cur.execute("CREATE INDEX mainindex on knowngene(chr,chr_start,chr_stop)")
	cur.execute("CREATE INDEX chrindex on knowngene(chr)")
	cur.execute("CREATE TABLE annotations  (id varchar,dbname varchar,annot varchar)")
        cur.execute("CREATE INDEX annotindex on annotations(id)")
        cur.execute("CREATE TABLE genelength ( id text, genelength int, exons int);")
	con.commit()

def add_entry(id, chr, strand, chr_start, chr_stop, type):
	cur.execute("INSERT INTO knowngene values (?,?,?,?,?,?)", (id,chr,strand,int(chr_start),int(chr_stop),type))


def add_annotation(id, db, annotation):
        cur.execute("INSERT INTO annotations values (?,?,?)", (id,db,annotation))


def add_gene_length(id, startexons, stopexons):
        gl=0
        for i in range(len(startexons)):
                sss=int(startexons[i])+1
                gl = gl + int(stopexons[i]) - int(startexons[i])
        con.execute("INSERT INTO genelength VALUES (?,?,?)", (id, gl, len(startexons)) )


def load_annotations_from_file(filename):
	f = open(filename)
	sys.stderr.write("Loading annotations from " + filename + "...\n")
	line = f.readline()
	i = 0
	counter=0
	while line != "":
        	id,annot = line.rstrip().split("\t")
	        add_annotation(id,os.path.splitext(os.path.basename(filename))[0],annot)
        	i = i + 1
	        if i % 1000 == 0:
        	        sys.stderr.write(str(i) + " annotations loaded...\n")
		if counter==300:
			con.commit()
			counter=-1
		counter+=1
	        line = f.readline()
	sys.stderr.write(str(i) + " total annotations loaded.\n")
	con.commit()

# db line : id,chr,strand,int(chr_start),int(chr_stop),type
# gtf line: chr,?,type,start,stop,?,strand,boh,string
def importGTF( f ):
	count=0
	for line in f :
		ff = line.strip().split("\t")
		gene=ff[-1].split("\s")
		myid=""
		old=""
		for aa in gene :
			if old == "gene_name" :
				myid=aa
			old=aa
		add_entry(myid,ff[0],ff[6], int(ff[3])-1, ff[4], ff[2].upper())
		count=count+1
		if count>300:
			count=0
			con.commit()
	con.commit()
	con.close()


### START MAIN PROGRAM
### START MAIN PROGRAM
### START MAIN PROGRAM

f = open(sys.argv[1])
initdb()

if "gtf" in sys.argv[1] :
	importGTF( f )
	sys.exit()

line = f.readline()
n = 0
counter=0
while line != "":
	splitted = line.split("\t")
	id = splitted[0]
	chr = splitted[1]
	strand = splitted[2]
	chr_start = int(splitted[3])+1
	chr_stop = splitted[4]
	cds_start = int(splitted[5])+1
	cds_stop = splitted[6]
	starts = splitted[8].split(",")
	stops = splitted[9].split(",")
	# insert introns
	if len(starts) > 1:
		for aa in range(0, len(starts)-1):
			if starts[aa+1] == '':
				continue
			iss = int(stops[aa]) + 1
			iee = int(starts[aa+1])
			add_entry(id,chr,strand,iss,iee,"INTRON")
	# adjust starts
	newstarts = []
	for start in starts:
		if start != "":
			newstarts.append(int(start)+1)
	starts = newstarts
        while len(starts)<len(stops):
                stops.pop()
        # update gene length table
        add_gene_length(id, starts, stops)

	add_entry(id,chr,strand,chr_start,chr_stop,"GENE")
	for i in range(len(starts)):
		if starts[i] != "":
			if int(cds_start) == int(cds_stop)+1:
				add_entry(id,chr,strand,starts[i],stops[i],"EXON")
			else:
				if int(starts[i]) < int(cds_start):
					# sono in un UTR
					if int(stops[i]) < int(cds_start):
						# se lo stop dell'esone e' prima dell'inizio del cds tutto l'esone e' nell'UTR
						add_entry(id,chr,strand,starts[i],stops[i],"UTR")
					else:
						# un pezzo nell'UTR e un pezzo in un CDS
						if int(stops[i]) <= int(cds_stop):
							add_entry(id,chr,strand,starts[i],int(cds_start)-1,"UTR")
							add_entry(id,chr,strand,cds_start,stops[i],"CDS")
						else:
							add_entry(id,chr,strand,starts[i],int(cds_start)-1,"UTR")
							add_entry(id,chr,strand,cds_start,cds_stop,"CDS")
							add_entry(id,chr,strand,int(cds_stop)+1,stops[i],"UTR")
				else:
					# sono in un CDS a meno che non sia oltre il cds stop
					if int(starts[i]) < int(cds_stop):
						# sono in un cds
						if int(stops[i]) <= int(cds_stop):
							add_entry(id,chr,strand,starts[i],stops[i],"CDS")
						else:
							add_entry(id,chr,strand,starts[i],cds_stop,"CDS")
							add_entry(id,chr,strand,int(cds_stop)+1,stops[i],"UTR")
					else:
						add_entry(id,chr,strand,starts[i],stops[i],"UTR")
	n = n + 1
	if n%100 == 0:
		sys.stderr.write("Loaded " + str(n) + " annotations...\n")
	if counter==300:
		con.commit()
		counter=-1
	counter+=1
	line = f.readline()
con.commit()

for file in sys.argv[3:]:
	load_annotations_from_file(file)

con.commit()
con.close()

